import json
import os
from pathlib import Path
import random
from typing import Any

import numpy as np
import matplotlib.pyplot as plt

class sen12_tools:
    def __init__(
        self, dst_dir, list_file
    ) -> None:
        self.dst_dir = dst_dir
        self.file_list = self._read_list_file(list_file)

    def _read_list_file(self, list_file):
        jh = json_helper(json_dir=list_file)
        dict_ = jh.get_dict()
        pair_list = []
        for opt_, sar_ in zip(dict_["OPT"], dict_["SAR"]):
            pair_list.append([
                os.path.join(self.dst_dir, opt_), 
                os.path.join(self.dst_dir, sar_), 
            ])
        return pair_list

    def __len__(self):
        return len(self.file_list)

    def get_wkt(self, img_names) -> str:
        """
        image name to wkt
        
        Param
        -----
        img_names str or tuple just use the first elem if tuple
        """
        if (type(img_names) == tuple):
            img_names = img_names[0]
        
        splits_ = img_names.split(os.sep)
        img_wkt = os.path.splitext(splits_[-1])[0]
        modal_wkt = splits_[-2]
        wkt = img_wkt.replace(modal_wkt+'_', '')
        
        return wkt
        
    def get_patch(self, filename, transform=None, **kwargs):
        '''
        filename(文件名),transfrome --> img
        '''
        try:
            img = plt.imread(filename)
        except:
            return None

        if len(img.shape) == 2:
            img = np.expand_dims(img, axis=2)

        return transform(img) if transform is not None else img

def open_sen12_as_np(
    src_file, 
    *, 
    type='RGB', 
    modal='', 
    max_ : np.ndarray = np.array([256., 256., 256.]),
    min_ : np.ndarray = np.array([0., 0., 0.]),  
    scale : int = 255
) -> np.ndarray:
    if type not in ['RGB', 'NIR', 'SAR']:
        raise ValueError(f"type: {type} not in ['RGB', 'NIR', 'SAR'] !")
    # if type == 'RGB':
    #     if modal == 'PS-RGB':
    #         min_ = np.array([0, 0, 0])
    #         max_ = np.array([256., 256., 256.])
    #         scale = 256
    #     else:
    #         # 本分支有实际意义
    #         min_ = np.array([0, 0, 0])
    #         max_ = np.array([2048, 2048, 2048])
    #         scale = 256
    # elif type == 'NIR':
    #     min_ = np.array([0])
    #     max_ = np.array([2048])
    #     scale = 256
    # elif type == 'SAR':
    #     min_ = np.array([0])
    #     max_ = np.array([93., 92., 92., 92.])
    #     # max_ = np.array([256., 256., 256., 256.])
    #     scale = 256
    
    img_np = plt.imread(src_file)

    img_np = (img_np.clip(*(min_, max_)) - min_) / (max_ - min_)
    img_np = (img_np * scale).astype(np.uint8)

    return img_np

class json_helper:
    '''
    Example
    -----
    >>> j = json_helper(dict_=test_dict)
    >>> j.to_json(save_dir="E:/workspace/SOMatch/preprocess/test.json")
    '''
    def __init__(self, dict_: dict=None, json_dir=None) -> None:
        if dict_ == None and json_dir == None:
            raise ValueError("both dict_ and json_dir are None!")
        if json_dir != None:
            self.json_dir = json_dir
            _ = self.from_json()
        else:
            self.dict_ = dict_

    def to_json(self, *, save_dir = "") -> None:
        with open(save_dir, 'w') as fp:
            json.dump(self.dict_, fp, indent='\t', sort_keys=True)

    def from_json(self) -> dict:
        with open(self.json_dir, 'r') as fp:
            self.dict_ = json.load(fp)

        return self.dict_

    def get_dict(self) -> dict:
        return self.dict_

    def get_content(self) -> Any:
        '''
        跟get_dict相同，返回的json内容可能不止是dict
        '''
        return self.dict_

def dict_to_RGB_plotimg(img_dict, save_dir, *, 
    img_type='RGB', ch_list=None, color=[255, 0, 0], mark_type='.', 
    save_kw='', modal='', pproc_func=None
):
    '''
    根据城市对应的img2pt绘图

    Param:
    ----
    img_dict: img2dict字典
    save_dir: 图像保存的路径
    img_type: 绘制图像的类型('RGB'/'NIR'/'SAR')
    ch_list: 通道索引，可使用img_type对应的默认设置，其长度应为1/3
    color: 特征点标记的颜色，通道数要与ch_list保持一致
    mark_type: 特征点标记类型，与kpt2mark同名参数一致
    save_kw: 保存图像时的前缀
    modal: 模态名('SAR-Intensity'/'MS'/'PS-RGB'/)
    pproc_func: 预处理函数，为None则不处理
    '''
    
    def kpt2mark(shape, kpt, mark_type='.'):
        if mark_type == '.':
            mark_pt = kpt
        elif mark_type == 'square':
            radius = 5
            mark_num = 8 * radius
            mark_pt = kpt[:, np.newaxis, :]
            mark_pt = np.tile(mark_pt, (1, mark_num, 1))
            ends = np.array([
                [-radius, -radius],
                [radius, -radius],
                [radius, radius],
                [-radius, radius]
                ])
            addend = np.array(range(0, 2*radius), dtype=np.int32)
            addend = np.stack((addend, np.zeros_like(addend)), axis=1)
            point_shift = np.empty((0, 2))
            for i in ends:
                if i[0] * i[1] < 0:
                    tmp_shift = addend[:, 1::-1]
                else:
                    tmp_shift = addend
                if i[1] > 0:
                    tmp_shift = -tmp_shift
                tmp_shift = tmp_shift + i
                point_shift = np.append(point_shift, tmp_shift, axis=0)
            mark_pt = mark_pt + point_shift
            mark_pt = mark_pt.reshape(-1, mark_pt.shape[2])
            mark_pt = mark_pt.clip([0, 0], [shape[0]-1, shape[1]-1])

        return mark_pt

    if not os.path.exists(save_dir):
        os.mkdir(save_dir)

    if img_type == 'RGB':
        min_ = np.array([0., 0., 0.])
        max_ = np.array([1., 1., 1.])
        if ch_list == None:
            ch_list = [0, 1, 2]
        elif len(ch_list) !=1 and len(ch_list) != 3:
            raise ValueError(f"length of ch_list: {ch_list} should be 1 or 3")
        file_modal = 'MS' if modal == '' else modal
    # elif img_type == 'NIR':
    #     if ch_list == None:
    #         ch_list = [0]
    #     elif len(ch_list) !=1:
    #         raise ValueError(f"length of ch_list: {ch_list} should be 1")
    #     file_modal = 'MS'
    elif img_type == 'SAR':
        min_ = np.array([0.])
        max_ = np.array([1.])
        if ch_list == None:
            # ch_list = [0]
            ch_list = [0, 1, 2]
        elif len(ch_list) !=1 and len(ch_list) != 3:
            raise ValueError(f"length of ch_list: {ch_list} should be 1 or 3")
        file_modal = 'SAR-Intensity'

    for file_, points in img_dict.items():
        background = open_sen12_as_np(
            file_, type=img_type, modal=modal, min_=min_, max_=max_
        )
        if len(background.shape) == 2:
            background = np.stack((background, background, background), axis=2)
        background = background[:, :, ch_list]

        background = pproc_func(background) if pproc_func != None else background

        fig = plt.figure()
        plt.imshow(background)
        plt.axis('off')

        # 计算标记形状
        # kpt_coord = np.empty((0, 2), dtype=np.int32)
        for pt_attri in points.values():
            # kpt_coord = np.append(kpt_coord, 
            #     np.array(pt_attri["xy"]).reshape(1, -1), axis=0)
            cur_coord = pt_attri["xy"]
            kpt_coord = [
                cur_coord[0]-5, cur_coord[1]-5, 
                cur_coord[0]+5, cur_coord[1]+5
                ]
            annot = round(pt_attri["response"], 3)

            # if kpt_coord.shape[0] > 0:
            #     mark_coord = kpt2mark(
            #         background.shape[1::-1], kpt_coord, mark_type=mark_type)
            #     mark_coord = mark_coord.astype(np.int32)
            #     # if len(background.shape) == 2 and len(color.shape) > 1:
            #     #     background = np.expand_dims(background, axis=2)
            #     #     tmp = background
            #     #     for _ in range(color.shape - 1):
            #     #         background = np.concatenate((background, tmp), axis=2)
            #     background[mark_coord[:, 1], mark_coord[:, 0]] = color

            plt.gca().add_patch(
                plt.Rectangle((kpt_coord[0], kpt_coord[1]), 
                    kpt_coord[2] - kpt_coord[0],
                    kpt_coord[3] - kpt_coord[1], fill=False,
                    edgecolor='r', linewidth=1)
                )
            plt.text(
                kpt_coord[0], kpt_coord[1]-1, str(annot), 
                # backgroundcolor = color, 
                color = (1, 1, 1),  bbox=dict(boxstyle='round,pad=0', fc=color, lw=0,  alpha=0.7)
            )

        showimg_name = os.path.basename(file_)
        showimg_name = os.path.splitext(showimg_name)[0] + ".png"
        showimg_name = save_kw + showimg_name if save_kw != '' else showimg_name
        
        # plt.imsave(os.path.join(save_dir, showimg_name), background)
        fig.tight_layout()
        plt.savefig(os.path.join(save_dir, showimg_name))
        plt.close(fig)

def get_sen_filelist(dset_dir, *, relpath=True, save_dir=""):
    '''
    将sen12地全部图片写入json文件，并分别生成每个场景的json文件
    '''
    json_ = {'OPT': [], 'SAR': []}
    p = Path(dset_dir)

    for roi in p.iterdir():
        if not roi.is_dir():
            continue

        cur_json = {'OPT': [], 'SAR': []}

        for region in roi.iterdir():
            if not region.is_dir():
                continue

            
            for filename in region.iterdir():
                if not filename.is_file():
                    continue

                if '_s1_' in str(filename):
                    if relpath:
                        cur_json['SAR'].append(
                            os.path.relpath(filename, p)
                        )
                    else:
                        cur_json['SAR'].append(
                            str(filename)
                        )
                elif '_s2_' in str(filename):
                    if relpath:
                        cur_json['OPT'].append(
                            os.path.relpath(filename, p)
                        )
                    else:
                        cur_json['OPT'].append(
                            str(filename)
                        )

        jh = json_helper(dict_=cur_json)
        jh.to_json(save_dir=os.path.join(save_dir, f"{str(os.path.basename(roi))}.json"))
        json_.update(cur_json)

    jh = json_helper(dict_=json_)
    jh.to_json(save_dir=os.path.join(save_dir, "sen12_total.json"))

def get_sen12_pairs(filename : str, *, dset_dir='') -> tuple:
    '''
    根据某一模态的图片名获得两种模态的图片名(s1, s2)，如果dset_dir不为''，则返回的图片路径为相对于该值的相对路径
    '''
    if (dset_dir != '' and dset_dir in filename):
        filename = os.path.relpath(filename, dset_dir)
    new_filename = None
    ret = None
    if '_s1_' in filename:
        new_filename = filename.replace('_s1_', '_s2_')
        new_filename = new_filename.replace(
            os.sep+'s1', os.sep+'s2')
        ret = (filename, new_filename)
    elif '_s2_' in filename:
        new_filename = filename.replace('_s2_', '_s1_')
        new_filename = new_filename.replace(
            os.sep+'s2', os.sep+'s1')
        ret = (new_filename, filename)

    return ret
    
def select_sen12_transaction(dset_dir, sar_json, opt_json, *, center=0, save_path='', roi='', threshold=None):
    '''
    根据img2pt信息筛选特征点在图片中央的图像

    Params
    -----
    dset_dir 数据集根目录，用于计算相对路径
    sar_json sar图像的img2pt路径
    opt_json opt图像的img2pt路径
    center   特征点两坐标离中心点的最大距离
    save_path 存储结果位置，格式{'OPT': [], 'SAR': []}的json
    threshold 响应大于阈值的点{}
    '''
    sar_jh = json_helper(json_dir=sar_json)
    opt_jh = json_helper(json_dir=opt_json)
    sar_dict = sar_jh.get_dict()
    opt_dict = opt_jh.get_dict()

    dict_ = {'OPT': [], 'SAR': []}

    # 在SAR图像中寻找
    for filename in sar_dict:
        if roi != '' and not roi in filename:
            continue
        for pt in sar_dict[filename].values():
            flag = False # 当前图片是否符合要求

            if threshold != None and pt["response"] < threshold['SAR']:
                continue

            x, y = pt['xy']
            if x - 128 > -center and x - 128 < center \
            and y - 128 > -center and y - 128 < center:
                pair = get_sen12_pairs(filename, dset_dir=dset_dir)
                if pair == None:
                    continue
                opt_key = os.path.join(dset_dir, pair[1])

                # 再在OPT图像中寻找
                for opt_pt in opt_dict[opt_key].values():
                    if threshold != None and opt_pt["response"] < threshold['OPT']:
                        continue

                    opt_x, opt_y = opt_pt['xy'] 
                    if x - 128 > -center and x - 128 < center \
                    and y - 128 > -center and y - 128 < center:
                        flag = True
                        break

                if flag:
                    dict_['OPT'].append(pair[1])
                    dict_['SAR'].append(pair[0])
                    break

    save_jh = json_helper(dict_=dict_)
    save_jh.to_json(save_dir=save_path)

def get_img_union_from_img2dict(*args, ret_modal=1) -> list:
    '''
    输入由若干文件名列表组成地列表，输出他们在相同模态下并集的列表

    Params:
    -----
    ret_modal: 1、返回的文件名是s2的opt，0、返回的文件名是s1的sar
    '''
    ret = []
    for file_list in args:
        for filename in file_list:
            modal_filename = get_sen12_pairs(filename)[ret_modal]
            if modal_filename not in ret:
                ret.append(modal_filename)

    return ret

def get_sen12_subset(
    fold_list: list, *, rate: float=None, num: int=None
) -> dict:
    '''
    划分指定占比或者数量的子数据集

    Param:
    -----
    fold_list:表模态名的文件夹路径列表
    rate:子数据集占总数据集比例
    num:子数据集的数量，和rate参数仅一个有效

    Return:
    -----
    返回子数据集列表{'SAR': [], 'OPT': []}
    '''

    if rate != None:
        num = round(rate * len(fold_list))
    elif num == None:
        raise ValueError('both rate and num are None!')

    shuffle_ids = list(range(len(fold_list)))
    random.shuffle(shuffle_ids)
    subset_ids = shuffle_ids[:num]

    subset = {'SAR': [], 'OPT': []}
    for ids in subset_ids:
        pair = get_sen12_pairs(fold_list[ids])
        subset['SAR'].append(pair[0])
        subset['OPT'].append(pair[1])

    return subset

def get_subset_main (json_path, subset_path,*, num=None, rate=None):
    """
    json_path       : "存放场景图像文件列表文件的路径"
    subset_path     : "子数据集列表文件存储路径"
    
    """

    jh = json_helper(json_dir=json_path)
    dict_ = jh.get_dict()
    fold_list = get_img_union_from_img2dict(dict_['OPT'], dict_['SAR'])

    subset_dict = get_sen12_subset(fold_list, num=num, rate=rate) 
    jh_subset = json_helper(dict_=subset_dict)
    jh_subset.to_json(save_dir=subset_path)


if __name__ == "__main__1":
    #---------------------- generate list
    # dset_dir = "E:/datasets/sen1-2"
    # save_dir = "E:/workspace/SOMatch/tmp/json/sen12_list"
    # get_sen_filelist(dset_dir, save_dir=save_dir)
    
    #---------------------- get subset
    json_path = "E:/workspace/SOMatch/tmp/json/sen12_list/sen12_total.json"
    subset_path = "E:/workspace/SOMatch/tmp/json/sen12_list/overlap_subset.json"
    get_subset_main(json_path, subset_path, num=200)


    # # test open_sen12_as_np
    # img_np = open_sen12_as_np(
    #     "E:/datasets/sen1-2/rois2017_winter/s2_23/ROIs2017_winter_s2_23_p1.png", 
    #     max_=np.array([1., 1., 1.])
    #     )
    # plt.imsave("E:/workspace/SOMatch/preprocess/a.png", img_np)


    #----------- test dict_to_RGB_plotimg


#---------------------- get subset
if __name__ == "__main__1":
    json_path = "E:/workspace/SOMatch/tmp/json/sen12_tt_harris/tt_pt.json"
    subset_path = "E:/workspace/SOMatch/tmp/json/sen12_tt_harris/pt_s100.json"
    get_subset_main(json_path, subset_path, num=100)


#------- draw the point searched from sen12
if __name__ == "__main__1":
    # folder to save results
    save_dir = "E:/workspace/SOMatch/image/sen12-200"
    save_kw = ''
    # sar 和 opt 的特征点文件
    opt_json_dir = "E:/workspace/SOMatch/tmp/json/sen12_ol_pt/subset_PS-RGB_HARRIS_subset200.json"
    sar_json_dir = "E:/workspace/SOMatch/tmp/json/sen12_ol_pt/subset_SAR_HARRIS_subset200.json"

    if not os.path.exists(save_dir):
        os.mkdir(save_dir)

    jh = json_helper(json_dir=opt_json_dir)

    dict_to_RGB_plotimg(
        jh.get_content(), save_dir, 
        img_type='RGB', ch_list=None, color=[1, 0, 0], mark_type='square', 
        save_kw=save_kw, modal=''
    )

    import sys
    sys.path.append("E:/workspace/SOMatch")
    from utils.preprocess import bilatera_blur

    jh = json_helper(json_dir=sar_json_dir)

    dict_to_RGB_plotimg(
        jh.get_content(), save_dir, 
        img_type='SAR', ch_list=None, color=[1, 0, 0], mark_type='square', 
        save_kw=save_kw, modal='', pproc_func=bilatera_blur
    )

    # # test sen12_tools
    # st = sen12_tools("E://datasets//sen1-2//", "E://workspace//SOMatch//preprocess//test_sen_subset.json")
    # print(st.get_wkt("E://datasets//sen1-2//ROIs1158_spring//s1_0//ROIs1158_spring_s1_0_p28.png"))

    # # test get_sen_filelist
    # get_sen_filelist("E://datasets//sen1-2//", relpath=False)

    # # test select_sen12
    # select_sen12(
    #     "E://datasets//sen1-2//", 
    #     "E://workspace//SOMatch//preprocess//sen12_subset//subset_SAR_HARRIS.json", 
    #     "E://workspace//SOMatch//preprocess//sen12_subset//subset_PS-RGB_HARRIS.json", 
    #     center=256 / 8, 
    #     save_path="E://workspace//SOMatch//preprocess//sen12_subset//sen12_select_rois1868_summer.json", roi="rois1868_summer"
    # )

    # # test select_sen12_transaction
    # select_sen12_transaction(
    #     "E://datasets//sen1-2//", 
    #     "E://workspace//SOMatch//preprocess//sen12_subset//subset_SAR_HARRIS.json", 
    #     "E://workspace//SOMatch//preprocess//sen12_subset//subset_PS-RGB_HARRIS.json", 
    #     center=256 / 8, 
    #     save_path="E://workspace//SOMatch//preprocess//sen12_subset//sen12_trans_rois1868_t045_122.json", roi="rois1868_summer", 
    #     threshold={'OPT': 1.22, 'SAR': 0.45}
    # )

    # # test get_img_union_from_img2dict
    # json_path = "E:/workspace/SOMatch/preprocess/sen12_subset/sen12_transaction_rois1868_summer.json"
    # jh = json_helper(json_dir=json_path)
    # dict_ = jh.get_dict()
    # ret_dict = get_img_union_from_img2dict(dict_['OPT'], dict_['SAR'])
    # jh_ret = json_helper(dict_=ret_dict)
    # jh_ret.to_json(save_dir='E:/workspace/SOMatch/preprocess/sen12_subset/sen12_transaction_rois1868_summer_union.json')

    # # test get_sen12_subset
    # fold_list_path = 'E:/workspace/SOMatch/preprocess/sen12_subset/sen12_transaction_rois1868_summer_union.json'
    # subset_path = "E:/workspace/SOMatch/preprocess/sen12_transaction_rois1868_summer_subset200.json"
    # jh = json_helper(json_dir=fold_list_path)
    # fold_list = jh.get_content()
    
    # subset_dict = get_sen12_subset(fold_list, num=200) 
    # jh_subset = json_helper(dict_=subset_dict)
    # jh_subset.to_json(save_dir=subset_path)

